import numpy as np
import matplotlib
matplotlib.use('TkAgg')  # 兼容 PyCharm 绘图后端
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# 解决中文乱码
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


# =======================
# 距离函数：闵可夫斯基距离
# =======================
def minkowski_distance(a, b, p=2):
    return np.sum(np.abs(a - b) ** p) ** (1 / p)


# =======================
# 自定义 KNN 分类器
# =======================
class KNN:
    def __init__(self, k=3, label_num=None, p=2):
        self.k = k
        self.label_num = label_num
        self.p = p  # p=1 曼哈顿距离, p=2 欧氏距离

    def fit(self, x_train, y_train):
        self.x_train = x_train
        self.y_train = y_train
        if self.label_num is None:
            self.label_num = len(np.unique(y_train))

    def get_knn_indices(self, x):
        distances = [minkowski_distance(a, x, self.p) for a in self.x_train]
        knn_indices = np.argsort(distances)[:self.k]
        return knn_indices

    def get_label(self, x):
        knn_indices = self.get_knn_indices(x)
        label_count = np.zeros(self.label_num)
        for idx in knn_indices:
            label = int(self.y_train[idx])
            label_count[label] += 1
        return np.argmax(label_count)

    def predict(self, x_test):
        predictions = np.zeros(len(x_test), dtype=int)
        for i, x in enumerate(x_test):
            predictions[i] = self.get_label(x)
        return predictions


# =======================
# 主程序入口
# =======================
if __name__ == "__main__":
    # 使用 Iris 数据集
    iris = load_iris()
    X = iris.data[:, :2]  # 取前两个特征便于可视化
    y = iris.target

    # 数据标准化
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    # 实例化并训练模型
    knn = KNN(k=5, p=2)
    knn.fit(X_train, y_train)

    # 预测
    y_pred = knn.predict(X_test)

    # 计算精度
    accuracy = np.mean(y_pred == y_test)
    print(f"预测准确率: {accuracy * 100:.2f}%")

    # ============ 可视化 ============
    x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
    y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
    xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.02),
                         np.arange(y_min, y_max, 0.02))

    # 网格点预测
    grid_points = np.c_[xx.ravel(), yy.ravel()]
    Z = knn.predict(grid_points)
    Z = Z.reshape(xx.shape)

    # 绘图
    plt.figure(figsize=(8, 6))
    plt.contourf(xx, yy, Z, alpha=0.3, cmap=plt.cm.rainbow)
    plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test, edgecolors='k', cmap=plt.cm.rainbow)
    plt.title(f"KNN 分类结果 (k={knn.k}, p={knn.p}, 准确率={accuracy:.2f})")
    plt.xlabel("Feature 1")
    plt.ylabel("Feature 2")
    plt.show()
